package com.tsing.cedap.common.config;

import cn.hutool.extra.servlet.ServletUtil;
import com.tsing.cedap.common.annotations.Limit;
import com.tsing.cedap.common.exceptions.BusinessException;
import com.tsing.cedap.common.exceptions.CommonException;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.annotation.Order;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

@Aspect
@Order(1)
@Component
public class LimitAspect {

    private final RedisTemplate<Object, Object> redisTemplate;

    public LimitAspect(RedisTemplate<Object, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Pointcut("@annotation(com.tsing.cedap.common.annotations.Limit)")
    public void pointcut() {
    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        String key;
        switch (limitAnnotation.limitType()) {
            case KEY:
                key = limitAnnotation.key();
                break;
            case IP:
                key = ServletUtil.getClientIP(request);
                break;
            case KEYIP:
                key = limitAnnotation.key() + "_" + ServletUtil.getClientIP(request);
                break;
            case METHOD:
                key = method.getName();
                break;
            default:
                key = method.getName() + "_" + ServletUtil.getClientIP(request);
        }
        List<Object> keys = Collections.singletonList(limitAnnotation.prefix() + key);
        RedisScript<Boolean> redisScript = new DefaultRedisScript<>(buildLuaScript(), Boolean.class);
        if (redisTemplate.execute(redisScript, keys, limitAnnotation.count(), limitAnnotation.period())) {
            return joinPoint.proceed();
        } else {
            throw new BusinessException(CommonException.Proxy.SERVER_IS_TOO_BUSY);
        }
    }

    /**
     * lua限流脚本
     */
    private String buildLuaScript() {
        return "local num = redis.call('incr',KEYS[1])\n" +
                "if tonumber(num) == 1 then\n" +
                "redis.call('pexpire',KEYS[1],ARGV[2])\n" +
                "return true\n" +
                "elseif tonumber(num) > tonumber(ARGV[1]) then\n" +
                "return false\n" +
                "else\n" +
                "return true\n" +
                "end";
    }

}
